#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
增强版LSTM-CRF模型
在训练时融合序列LSTM预测作为算子，并使用模式记忆网络模式的预测结果作为权重因子
支持排列5、大乐透和双色球
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torchcrf import CRF
import numpy as np
from datetime import datetime
import os
import sys

# 设置环境变量以强制使用CPU
# os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

# 添加项目根目录到路径
current_dir = os.path.dirname(os.path.abspath(__file__))
project_root = os.path.abspath(os.path.join(current_dir, '..'))
sys.path.append(project_root)

# 根据彩票类型导入相应的数据处理器
from algorithms.pattern_memory import PatternMemoryNetwork

# 修复相对导入问题
try:
    from model import LstmCRFModel
except ImportError:
    # 如果直接运行此文件，尝试其他导入方式
    try:
        import importlib.util
        spec = importlib.util.spec_from_file_location("model", os.path.join(project_root, "model.py"))
        model_module = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(model_module)
        LstmCRFModel = model_module.LstmCRFModel
    except:
        # 创建一个简单的替代类
        class LstmCRFModel(nn.Module):
            def __init__(self, input_dim, hidden_dim, output_dim, output_seq_length, num_layers=1, dropout=0.1):
                super(LstmCRFModel, self).__init__()
                self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=num_layers, batch_first=True)
                self.fc = nn.Linear(hidden_dim, output_seq_length * output_dim)
                self.crf = CRF(output_dim, batch_first=True)

class EnhancedLstmCRFModel(nn.Module):
    """增强版LSTM-CRF模型"""
    
    def __init__(self, input_dim, hidden_dim, output_dim, output_seq_length, num_layers=1, dropout=0.1):
        super(EnhancedLstmCRFModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.output_dim = output_dim
        self.output_seq_length = output_seq_length
        
        # 基础LSTM层
        self.lstm = nn.LSTM(input_dim, hidden_dim, num_layers=num_layers, batch_first=True, dropout=dropout if num_layers > 1 else 0)
        
        # 增强模块 - 融合序列LSTM预测
        # 修复：根据彩票类型动态设置增强模块的输入维度
        if output_seq_length == 3:  # fc3d
            self.sequence_lstm_enhancer = nn.Linear(hidden_dim + 10, hidden_dim)  # +10 for sequence LSTM predictions
        elif output_seq_length == 5:  # plw
            self.sequence_lstm_enhancer = nn.Linear(hidden_dim + 10, hidden_dim)  # +10 for sequence LSTM predictions
        elif output_seq_length == 20:  # kl8
            self.sequence_lstm_enhancer = nn.Linear(hidden_dim + 81, hidden_dim)  # +81 for sequence LSTM predictions
        else:  # dlt/ssq
            self.sequence_lstm_enhancer = nn.Linear(hidden_dim + output_dim, hidden_dim)
        
        # 全连接层
        self.fc = nn.Linear(hidden_dim, output_seq_length * output_dim)
        self.dropout = nn.Dropout(dropout)
        self.crf = CRF(output_dim, batch_first=True)
        
    def forward(self, x, sequence_lstm_predictions=None, target=None, mask=None):
        """
        前向传播
        
        Args:
            x: 输入特征
            sequence_lstm_predictions: 序列LSTM预测结果作为增强因子
            target: 目标标签（训练时使用）
            mask: 序列掩码
        """
        try:
            # 确保输入数据在正确的设备上
            if hasattr(self, 'lstm'):
                device = next(self.lstm.parameters()).device
                if x.device != device:
                    x = x.to(device)
            
            # LSTM层
            lstm_out, _ = self.lstm(x)
            
            # 如果提供了序列LSTM预测结果，则进行增强
            if sequence_lstm_predictions is not None:
                try:
                    # 确保序列LSTM预测结果在正确的设备上
                    if sequence_lstm_predictions.device != lstm_out.device:
                        sequence_lstm_predictions = sequence_lstm_predictions.to(lstm_out.device)
                    
                    # 将序列LSTM预测结果与LSTM输出融合
                    batch_size, seq_len, hidden_dim = lstm_out.shape
                    
                    # 检查sequence_lstm_predictions的形状
                    if sequence_lstm_predictions.dim() == 2:
                        # 扩展序列LSTM预测结果以匹配LSTM输出的时间步
                        expanded_seq_preds = sequence_lstm_predictions.unsqueeze(1).expand(-1, seq_len, -1)
                    else:
                        # 如果已经是正确的形状，直接使用
                        expanded_seq_preds = sequence_lstm_predictions
                    
                    # 确保维度匹配后再融合
                    if lstm_out.shape[0] == expanded_seq_preds.shape[0]:
                        # 检查增强模块的输入维度是否匹配
                        actual_input_dim = lstm_out.shape[-1] + expanded_seq_preds.shape[-1]
                        expected_input_dim = self.sequence_lstm_enhancer.in_features
                        
                        if actual_input_dim == expected_input_dim:
                            # 融合
                            enhanced_input = torch.cat([lstm_out, expanded_seq_preds], dim=-1)
                            lstm_out = torch.relu(self.sequence_lstm_enhancer(enhanced_input))
                        else:
                            print(f"增强模块输入维度不匹配: 实际 {actual_input_dim}, 期望 {expected_input_dim}，跳过增强")
                except Exception as e:
                    # 如果增强过程出错，使用原始LSTM输出
                    print(f"增强模块处理出错: {e}，使用原始LSTM输出")
            
            # 应用dropout
            lstm_out = self.dropout(lstm_out)
            
            # 全连接层
            fc_out = self.fc(lstm_out[:, -1, :])
            fc_out = fc_out.view(-1, self.output_seq_length, self.output_dim)
            
            # 如果提供了目标和掩码，计算损失
            if target is not None:
                # 确保目标在正确的设备上
                if target.device != fc_out.device:
                    target = target.to(fc_out.device)
                
                # 确保目标和输出形状匹配
                if target.dim() == 2:
                    target = target.unsqueeze(-1)  # 添加维度以匹配CRF期望的形状
                elif target.dim() == 1:
                    target = target.unsqueeze(0).unsqueeze(-1)
                
                # 计算CRF损失
                try:
                    loss = -self.crf(fc_out, target.squeeze(-1), reduction='mean')
                    return loss  # 返回标量损失
                except Exception as crf_error:
                    print(f"CRF损失计算出错: {crf_error}")
                    return torch.tensor(0.0, device=fc_out.device)  # 返回零损失
            
            # 否则进行解码
            try:
                predictions = self.crf.decode(fc_out)
                return predictions
            except Exception as decode_error:
                print(f"CRF解码出错: {decode_error}")
                # 返回默认预测结果
                return [[0] * self.output_seq_length]  # 返回默认预测
        except Exception as e:
            print(f"前向传播出错: {e}")
            import traceback
            print(traceback.format_exc())
            # 返回默认预测结果
            if target is not None:
                # 确保返回的张量在正确的设备上
                device = next(self.parameters()).device if next(self.parameters()).is_cuda else torch.device('cpu')
                return torch.tensor(0.0, device=device)  # 返回标量损失
            else:
                # 返回默认预测结果
                return [[0] * self.output_seq_length]  # 返回默认预测

class EnhancedLstmCRFTrainer:
    """增强版LSTM-CRF模型训练器"""
    
    def __init__(self, csv_file_path, model_save_path=None, lottery_type='plw'):
        self.csv_file_path = csv_file_path
        self.model_save_path = model_save_path or "enhanced_lstm_crf_model.pth"
        self.lottery_type = lottery_type
        
        # 初始化组件
        self.pattern_memory = PatternMemoryNetwork(lottery_type=lottery_type)
        self.sequence_lstm_model = None
        self.enhanced_model = None
        
        # 根据彩票类型设置参数
        if lottery_type == 'plw':
            self.window_size = 10
            self.input_dim = 10  # 修复：排列5应该使用10个特征（5个原始数字+5个区域转换特征）
            self.output_dim = 10
            self.output_seq_length = 5
        elif lottery_type == 'dlt':
            self.window_size = 10
            self.input_dim = 7  # 5个红球 + 2个蓝球
            self.output_dim = 35  # 大乐透红球范围1-35
            self.output_seq_length = 5  # 预测5个红球
        elif lottery_type == 'ssq':
            self.window_size = 10
            self.input_dim = 7  # 6个红球 + 1个蓝球
            self.output_dim = 33  # 双色球红球范围1-33
            self.output_seq_length = 6  # 预测6个红球
        elif lottery_type == 'kl8':
            self.window_size = 10
            self.input_dim = 40  # 20个号码，每个号码2个特征（号码+权重）
            self.output_dim = 81  # 快乐8号码范围1-80，加上0作为填充/特殊标记
            self.output_seq_length = 20  # 预测20个号码
        elif lottery_type == 'fc3d':
            self.window_size = 10
            self.input_dim = 3  # 3D彩票有3个数字位置
            self.output_dim = 10  # 数字范围0-9
            self.output_seq_length = 3  # 预测3个数字
        else:
            raise ValueError(f"不支持的彩票类型: {lottery_type}")
    
    def load_and_prepare_data(self):
        """加载和准备训练数据"""
        print(f"[LOAD] 加载和准备{self.lottery_type}训练数据...")
        
        # 根据彩票类型使用相应的数据处理器
        if self.lottery_type == 'plw':
            from algorithms.plw_sequence_lstm import PLWDataProcessor
            processor = PLWDataProcessor(self.csv_file_path, self.window_size)
            X, y = processor.load_and_process_data()
            # 为增强版LSTM-CRF模型准备正确的输入数据（使用完整的10个特征，包括区域转换特征）
            X_enhanced = X  # 使用完整的特征，包括区域转换特征
        elif self.lottery_type == 'kl8':
            # 快乐8使用自定义的数据处理器
            from scripts.kl8.train_kl8_model import prepare_training_data
            import pandas as pd
            import numpy as np
            
            # 读取快乐8历史数据
            df = pd.read_csv(self.csv_file_path, encoding='utf-8')
            # 提取开奖号码（第3列到第22列，对应1号到20号）
            numbers = df.iloc[:, 2:22].values.astype(int)
            
            # 准备训练数据
            X, y = prepare_training_data(numbers, self.window_size)
            X_enhanced = X.reshape(X.shape[0], X.shape[1], -1)  # 展平特征维度
        elif self.lottery_type == 'fc3d':
            # 3D彩票使用自定义的数据处理器
            from scripts.fc3d.train_3d_model import D3Dataset
            import pandas as pd
            import numpy as np
            
            # 读取3D历史数据
            dataset = D3Dataset(self.csv_file_path, self.window_size)
            X, y = dataset.features, dataset.labels
            X_enhanced = X  # 使用完整的特征
        else:
            # 对于大乐透和双色球，使用通用的数据处理器
            if self.lottery_type == 'dlt':
                from scripts.dlt.train_dlt_model import LotteryDataset
                dataset = LotteryDataset(self.csv_file_path, self.window_size, red_balls=5, blue_balls=2)
            else:  # ssq
                from scripts.ssq.train_ssq_model import LotteryDataset
                dataset = LotteryDataset(self.csv_file_path, self.window_size, red_balls=6, blue_balls=1)
            
            X, y = dataset.features, dataset.labels
            X_enhanced = X  # 使用完整的特征
        
        print(f"原始数据形状: X={X.shape}, y={y.shape}")
        print(f"增强版LSTM-CRF模型输入数据形状: X={X_enhanced.shape}, y={y.shape}")
        return X_enhanced, y
    
    def train_pattern_memory(self, historical_data):
        """训练模式记忆网络"""
        print("[TRAIN] 训练模式记忆网络...")
        self.pattern_memory.train(historical_data)
        print("[SUCCESS] 模式记忆网络训练完成")
    
    def get_pattern_memory_weights(self, recent_data):
        """获取模式记忆网络的预测权重"""
        try:
            # 使用模式记忆网络生成预测
            pattern_pred = self.pattern_memory.predict_single_set(recent_data)
            confidence = pattern_pred.get('confidence', 0.5)
            return confidence
        except Exception as e:
            print(f" 获取模式记忆权重失败: {e}")
            return 0.5  # 默认权重
    
    def load_sequence_lstm_model(self, model_path):
        """加载序列LSTM模型"""
        try:
            if os.path.exists(model_path):
                print(f"[LOAD] 加载序列LSTM模型: {model_path}")
                device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
                checkpoint = torch.load(model_path, map_location=device)
                
                # 根据彩票类型创建相应的序列LSTM模型
                if self.lottery_type == 'plw':
                    from algorithms.plw_sequence_lstm import PLWSequenceLSTM
                    self.sequence_lstm_model = PLWSequenceLSTM(
                        input_dim=10,  # 修复参数名称：input_size -> input_dim
                        hidden_dim=128,
                        num_layers=3,
                        dropout=0.3
                    )
                elif self.lottery_type == 'kl8':
                    # 快乐8使用自定义的LSTM模型
                    from scripts.kl8.train_kl8_model import KL8LSTMCRF
                    self.sequence_lstm_model = KL8LSTMCRF(
                        input_size=40,  # 20个号码，每个号码2个特征
                        hidden_size=128,
                        num_layers=2,
                        num_tags=81,  # 号码范围1-80，加上0作为填充/特殊标记
                        dropout=0.2
                    )
                elif self.lottery_type == 'fc3d':
                    # 3D彩票使用通用的LSTM模型
                    from model import LstmCRFModel
                    self.sequence_lstm_model = LstmCRFModel(
                        input_dim=3,  # 3D彩票有3个数字位置
                        hidden_dim=128,
                        output_dim=10,  # 数字范围0-9
                        output_seq_length=3,  # 预测3个数字
                        num_layers=2,
                        dropout=0.2
                    )
                else:
                    # 对于大乐透和双色球，使用通用的LSTM模型
                    from model import LstmCRFModel
                    # 根据彩票类型设置模型参数
                    if self.lottery_type == 'dlt':
                        input_dim = 7  # 5个红球 + 2个蓝球
                        output_dim = 35  # 大乐透红球范围
                        output_seq_length = 5
                    else:  # ssq
                        input_dim = 7  # 6个红球 + 1个蓝球
                        output_dim = 33  # 双色球红球范围
                        output_seq_length = 6
                    
                    self.sequence_lstm_model = LstmCRFModel(
                        input_dim=input_dim,
                        hidden_dim=128,
                        output_dim=output_dim,
                        output_seq_length=output_seq_length,
                        num_layers=2,
                        dropout=0.2
                    )
                
                # 检查检查点的格式并加载模型
                if isinstance(checkpoint, dict):
                    # 检查常见的键
                    if 'model_state_dict' in checkpoint:
                        self.sequence_lstm_model.load_state_dict(checkpoint['model_state_dict'])
                    elif 'state_dict' in checkpoint:
                        # 有些模型可能使用'state_dict'作为键
                        self.sequence_lstm_model.load_state_dict(checkpoint['state_dict'])
                    elif 'model' in checkpoint:
                        # fc3d的单独模型保存格式
                        self.sequence_lstm_model.load_state_dict(checkpoint['model'])
                    elif 'lstm_model' in checkpoint:
                        # fc3d的集成模型格式
                        self.sequence_lstm_model.load_state_dict(checkpoint['lstm_model'])
                    elif 'red_model' in checkpoint and 'blue_model' in checkpoint:
                        # 大乐透和双色球的特殊格式（红球和蓝球模型）
                        # 这里我们只需要加载红球模型用于增强版LSTM-CRF
                        self.sequence_lstm_model.load_state_dict(checkpoint['red_model'])
                    elif 'red_lstm_model' in checkpoint:
                        # 集成模型格式
                        self.sequence_lstm_model.load_state_dict(checkpoint['red_lstm_model'])
                    else:
                        # 如果检查点本身就是状态字典
                        try:
                            self.sequence_lstm_model.load_state_dict(checkpoint)
                        except Exception as e:
                            print(f"无法加载模型状态字典: {e}")
                            return False
                else:
                    # 检查点不是字典，尝试直接加载
                    try:
                        self.sequence_lstm_model.load_state_dict(checkpoint)
                    except Exception as e:
                        print(f"无法加载模型状态字典: {e}")
                        return False
                
                self.sequence_lstm_model.to(device)
                self.sequence_lstm_model.eval()
                
                print("[SUCCESS] 序列LSTM模型加载成功")
                return True
            else:
                print(f"[ERROR] 序列LSTM模型文件不存在: {model_path}")
                return False
        except Exception as e:
            print(f"[ERROR] 序列LSTM模型加载失败: {e}")
            import traceback
            print(traceback.format_exc())
            return False
    
    def get_sequence_lstm_predictions(self, input_data):
        """获取序列LSTM预测结果"""
        if self.sequence_lstm_model is None:
            print("[WARN] 序列LSTM模型未加载")
            return None
            
        try:
            # 确保输入数据和模型在同一设备上
            device = next(self.sequence_lstm_model.parameters()).device
            if input_data.device != device:
                input_data = input_data.to(device)
            
            with torch.no_grad():
                if self.lottery_type == 'plw':
                    from algorithms.plw_sequence_lstm import PLWSequenceLSTM
                    predictions, probabilities = self.sequence_lstm_model.predict(input_data)
                elif self.lottery_type == 'kl8':
                    # 快乐8使用自定义模型预测
                    predictions = self.sequence_lstm_model(input_data)
                elif self.lottery_type == 'fc3d':
                    # 3D彩票使用通用模型预测
                    predictions = self.sequence_lstm_model(input_data)
                else:
                    # 对于大乐透和双色球，直接使用模型预测
                    predictions = self.sequence_lstm_model(input_data)
                
                # 添加类型和形状检查
                if predictions is not None:
                    # 如果predictions是列表，尝试转换为张量
                    if isinstance(predictions, list):
                        print(f"序列LSTM预测结果是列表，长度: {len(predictions)}")
                        # 如果列表中的元素是张量，合并它们
                        if all(isinstance(p, torch.Tensor) for p in predictions):
                            # 尝试将列表中的张量合并
                            try:
                                predictions = torch.stack(predictions)
                            except Exception as e:
                                print(f"无法将列表中的张量合并: {e}")
                                # 如果无法合并，返回第一个元素（如果存在）
                                if len(predictions) > 0:
                                    predictions = predictions[0]
                                else:
                                    predictions = None
                        else:
                            # 如果列表中的元素不是张量，尝试转换
                            try:
                                predictions = torch.tensor(predictions)
                            except Exception as e:
                                print(f"无法将列表转换为张量: {e}")
                                predictions = None
                    elif isinstance(predictions, torch.Tensor):
                        print(f"序列LSTM预测结果形状: {predictions.shape}")
                    else:
                        print(f"序列LSTM预测结果类型: {type(predictions)}")
                    
                    # 修复：确保返回的预测结果在正确的设备上
                    if isinstance(predictions, torch.Tensor):
                        predictions = predictions.to(device)

                    return predictions
                else:
                    print("[WARN] 序列LSTM预测结果为空")
                    return None
        except Exception as e:
            print(f" 序列LSTM预测失败: {e}")
            import traceback
            print(traceback.format_exc())
            return None
    
    def create_enhanced_model(self, input_dim=None, hidden_dim=128, output_dim=None, output_seq_length=None):
        """创建增强版LSTM-CRF模型"""
        if input_dim is None:
            input_dim = self.input_dim
        if output_dim is None:
            output_dim = self.output_dim
        if output_seq_length is None:
            output_seq_length = self.output_seq_length
            
        print(f"[CREATE] 创建{self.lottery_type}增强版LSTM-CRF模型...")
        self.enhanced_model = EnhancedLstmCRFModel(
            input_dim=input_dim,
            hidden_dim=hidden_dim,
            output_dim=output_dim,
            output_seq_length=output_seq_length,
            num_layers=2,
            dropout=0.2
        )
        print("[SUCCESS] 增强版LSTM-CRF模型创建完成")
        return self.enhanced_model

    def train_model(self, X, y, epochs=100, batch_size=32, learning_rate=0.001):
        """训练增强版LSTM-CRF模型"""
        if self.enhanced_model is None:
            raise ValueError("模型未创建，请先调用create_enhanced_model方法")
        
        print(f"[RUN] 开始训练{self.lottery_type}增强版LSTM-CRF模型...")
        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        self.enhanced_model.to(device)
        
        # 转换数据为PyTorch张量并确保在同一设备上
        X_tensor = torch.FloatTensor(X).to(device)
        
        # 根据彩票类型处理标签
        if self.lottery_type == 'plw':
            # 排列5: 使用完整的5个数字标签
            y_tensor = torch.LongTensor(y).to(device)
        elif self.lottery_type == 'kl8':
            # 快乐8: 使用完整的20个数字标签
            y_tensor = torch.LongTensor(y).to(device)
        elif self.lottery_type == 'fc3d':
            # 3D彩票: 使用完整的3个数字标签
            y_tensor = torch.LongTensor(y).to(device)
        else:
            # 大乐透和双色球: 只使用红球部分的标签
            if self.lottery_type == 'dlt':
                # 大乐透: 只取前5个红球标签
                y_tensor = torch.LongTensor(y[:, :5]).to(device)
            else:  # ssq
                # 双色球: 只取前6个红球标签
                y_tensor = torch.LongTensor(y[:, :6]).to(device)
        
        # 创建数据加载器
        from torch.utils.data import DataLoader, TensorDataset
        dataset = TensorDataset(X_tensor, y_tensor)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
        
        # 定义优化器和损失函数
        optimizer = optim.AdamW(self.enhanced_model.parameters(), lr=learning_rate, weight_decay=1e-4)
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=10)
        
        # 训练循环
        best_loss = float('inf')
        patience_counter = 0
        patience_limit = 20
        
        for epoch in range(epochs):
            self.enhanced_model.train()
            total_loss = 0
            num_batches = 0
            
            for batch_idx, (batch_X, batch_y) in enumerate(dataloader):
                try:
                    optimizer.zero_grad()
                    
                    # 获取序列LSTM预测结果作为增强因子
                    sequence_lstm_preds = self.get_sequence_lstm_predictions(batch_X)
                    
                    sequence_lstm_onehot = None
                    if sequence_lstm_preds is not None:
                        try:
                            # 检查预测结果的形状和类型
                            print(f" 批次 {batch_idx}: 序列LSTM预测结果类型: {type(sequence_lstm_preds)}, 形状: {getattr(sequence_lstm_preds, 'shape', 'N/A')}")
                            
                            # 根据彩票类型转换为one-hot编码
                            batch_size_actual = batch_X.shape[0]
                            
                            if self.lottery_type == 'fc3d':
                                # 3D彩票: 3个数字位置，每个数字0-9
                                sequence_lstm_onehot = torch.zeros(batch_size_actual, 10).to(device)
                                
                                # 确保sequence_lstm_preds是正确的类型和形状
                                if isinstance(sequence_lstm_preds, torch.Tensor):
                                    if sequence_lstm_preds.dim() == 2:  # [batch_size, seq_length]
                                        for i in range(batch_size_actual):
                                            for pos in range(min(3, sequence_lstm_preds.shape[1])):  # 3D彩票只有3个位置
                                                digit = sequence_lstm_preds[i, pos].item()
                                                if 0 <= digit < 10:
                                                    sequence_lstm_onehot[i, digit] += 1
                                    elif sequence_lstm_preds.dim() == 1:  # [batch_size]
                                        for i in range(batch_size_actual):
                                            digit = sequence_lstm_preds[i].item()
                                            if 0 <= digit < 10:
                                                sequence_lstm_onehot[i, digit] += 1
                            elif self.lottery_type == 'plw':
                                # 排列5: 5个数字位置，每个数字0-9
                                sequence_lstm_onehot = torch.zeros(batch_size_actual, 10).to(device)
                                
                                if isinstance(sequence_lstm_preds, torch.Tensor):
                                    if sequence_lstm_preds.dim() == 2:  # [batch_size, seq_length]
                                        for i in range(batch_size_actual):
                                            for pos in range(min(5, sequence_lstm_preds.shape[1])):  # 排列5有5个位置
                                                digit = sequence_lstm_preds[i, pos].item()
                                                if 0 <= digit < 10:
                                                    sequence_lstm_onehot[i, digit] += 1
                                    elif sequence_lstm_preds.dim() == 1:  # [batch_size]
                                        for i in range(batch_size_actual):
                                            digit = sequence_lstm_preds[i].item()
                                            if 0 <= digit < 10:
                                                sequence_lstm_onehot[i, digit] += 1
                            elif self.lottery_type == 'kl8':
                                # 快乐8: 20个数字位置，每个数字1-80
                                sequence_lstm_onehot = torch.zeros(batch_size_actual, 81).to(device)  # 0-80，0作为填充
                                
                                if isinstance(sequence_lstm_preds, torch.Tensor):
                                    if sequence_lstm_preds.dim() == 2:  # [batch_size, seq_length]
                                        for i in range(batch_size_actual):
                                            for pos in range(min(20, sequence_lstm_preds.shape[1])):  # 快乐8有20个位置
                                                digit = sequence_lstm_preds[i, pos].item()
                                                if 0 <= digit <= 80:  # 快乐8号码范围1-80，0作为特殊值
                                                    sequence_lstm_onehot[i, digit] += 1
                                    elif sequence_lstm_preds.dim() == 1:  # [batch_size]
                                        for i in range(batch_size_actual):
                                            digit = sequence_lstm_preds[i].item()
                                            if 0 <= digit <= 80:
                                                sequence_lstm_onehot[i, digit] += 1
                            else:
                                # 其他类型使用默认处理
                                sequence_lstm_onehot = torch.zeros(batch_size_actual, 10).to(device)
                                
                                if isinstance(sequence_lstm_preds, torch.Tensor):
                                    if sequence_lstm_preds.dim() == 2:  # [batch_size, seq_length]
                                        for i in range(batch_size_actual):
                                            for pos in range(min(5, sequence_lstm_preds.shape[1])):
                                                digit = sequence_lstm_preds[i, pos].item()
                                                if 0 <= digit < 10:
                                                    sequence_lstm_onehot[i, digit] += 1
                                    elif sequence_lstm_preds.dim() == 1:  # [batch_size]
                                        for i in range(batch_size_actual):
                                            digit = sequence_lstm_preds[i].item()
                                            if 0 <= digit < 10:
                                                sequence_lstm_onehot[i, digit] += 1
                        except Exception as e:
                            print(f"处理序列LSTM预测结果时出错: {e}")
                            import traceback
                            print(traceback.format_exc())
                            sequence_lstm_onehot = None
                    else:
                        print(f" 批次 {batch_idx}: 序列LSTM预测结果为空")
                    
                    # 前向传播
                    try:
                        loss = self.enhanced_model(
                            batch_X, 
                            sequence_lstm_predictions=sequence_lstm_onehot,
                            target=batch_y
                        )
                        
                        # 确保损失是标量
                        if isinstance(loss, (list, tuple)):
                            loss = loss[0]  # 取第一个元素
                        
                        # 反向传播
                        loss.backward()
                        torch.nn.utils.clip_grad_norm_(self.enhanced_model.parameters(), max_norm=1.0)
                        optimizer.step()
                        
                        total_loss += loss.item()
                        num_batches += 1
                        
                        # 添加批次处理完成的提示
                        if batch_idx % 10 == 0:
                            print(f"[PROCESS] 批次 {batch_idx} 处理完成")
                    except Exception as forward_error:
                        print(f"[ERROR] 前向传播或反向传播出错 (批次 {batch_idx}): {forward_error}")
                        import traceback
                        print(traceback.format_exc())
                        # 继续处理下一个批次而不是中断整个训练
                        continue
                        
                except Exception as e:
                    print(f"[ERROR] 批次 {batch_idx} 处理出错: {e}")
                    import traceback
                    print(traceback.format_exc())
                    # 继续处理下一个批次而不是中断整个训练
                    continue
            
            avg_loss = total_loss / num_batches if num_batches > 0 else 0
            scheduler.step(avg_loss)
            
            # 打印进度
            if (epoch + 1) % 10 == 0:
                print(f"[EPOCH] Epoch {epoch+1:3d}/{epochs}: Loss = {avg_loss:.4f}")
            
            # 保存最佳模型
            if avg_loss < best_loss:
                best_loss = avg_loss
                patience_counter = 0
                self.save_model()
                print(f"[SAVE] 保存最佳模型 (Epoch {epoch+1}, Loss: {avg_loss:.4f})")
            else:
                patience_counter += 1
                if patience_counter >= patience_limit:
                    print(f"[STOP] 早停触发 (Patience: {patience_limit})")
                    break
        
        print(f"[SUCCESS] 训练完成！最佳损失: {best_loss:.4f}")
        return best_loss
    
    def save_model(self):
        """保存模型"""
        if self.enhanced_model is None:
            return
            
        try:
            torch.save({
                'model_state_dict': self.enhanced_model.state_dict(),
                'model_config': {
                    'input_dim': self.input_dim,  # 使用训练器中设置的正确输入维度
                    'hidden_dim': self.enhanced_model.hidden_dim,
                    'output_dim': self.enhanced_model.output_dim,
                    'output_seq_length': self.enhanced_model.output_seq_length,
                    'num_layers': self.enhanced_model.lstm.num_layers,  # 添加层数信息
                    'dropout': self.enhanced_model.dropout.p  # 添加dropout信息
                }
            }, self.model_save_path)
            print(f"[SAVE] 模型保存到: {self.model_save_path}")
        except Exception as e:
            print(f"  模型保存失败: {e}")
    
    def load_model(self):
        """加载模型"""
        try:
            if os.path.exists(self.model_save_path):
                device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
                checkpoint = torch.load(self.model_save_path, map_location=device)
                
                config = checkpoint['model_config']
                self.enhanced_model = EnhancedLstmCRFModel(**config)
                self.enhanced_model.load_state_dict(checkpoint['model_state_dict'])
                self.enhanced_model.to(device)
                self.enhanced_model.eval()
                
                print("[SUCCESS] 增强版LSTM-CRF模型加载成功")
                return True
            else:
                print(f" 模型文件不存在: {self.model_save_path}")
                return False
        except Exception as e:
            print(f"  模型加载失败: {e}")
            return False
    
    def predict(self, input_data):
        """使用增强版模型进行预测"""
        if self.enhanced_model is None:
            raise ValueError("模型未加载，请先调用load_model方法")
        
        # 确保输入数据是torch.Tensor类型
        if not isinstance(input_data, torch.Tensor):
            # 先创建CPU张量
            input_tensor = torch.FloatTensor(input_data)
        else:
            input_tensor = input_data
            
        # 获取模型所在的设备
        device = next(self.enhanced_model.parameters()).device
        # 将输入数据移动到正确的设备
        input_tensor = input_tensor.to(device)
        
        # 获取序列LSTM预测结果作为增强因子
        sequence_lstm_preds = self.get_sequence_lstm_predictions(input_tensor)
        
        sequence_lstm_onehot = None
        if sequence_lstm_preds is not None:
            try:
                # 转换为one-hot编码
                batch_size_actual = input_tensor.shape[0]
                sequence_lstm_onehot = torch.zeros(batch_size_actual, 10).to(device)
                
                # 确保sequence_lstm_preds是正确的类型和形状
                if isinstance(sequence_lstm_preds, torch.Tensor):
                    if sequence_lstm_preds.dim() == 2:  # [batch_size, seq_length]
                        for i in range(batch_size_actual):
                            for pos in range(min(5, sequence_lstm_preds.shape[1])):
                                digit = sequence_lstm_preds[i, pos].item()
                                if 0 <= digit < 10:
                                    sequence_lstm_onehot[i, digit] += 1
                    elif sequence_lstm_preds.dim() == 1:  # [batch_size]
                        for i in range(batch_size_actual):
                            digit = sequence_lstm_preds[i].item()
                            if 0 <= digit < 10:
                                sequence_lstm_onehot[i, digit] += 1
            except Exception as e:
                print(f" 处理序列LSTM预测结果时出错: {e}")
                sequence_lstm_onehot = None
        else:
            print(" 序列LSTM预测结果为空")
        
        # 进行预测
        with torch.no_grad():
            predictions = self.enhanced_model(input_tensor, sequence_lstm_predictions=sequence_lstm_onehot)
            # 修复：确保返回有效的预测结果
            if predictions is None:
                # 如果预测结果为空，返回默认值
                print(" 警告：模型预测结果为空，返回默认预测")
                # 返回一个默认的预测结果
                default_predictions = [[0, 1, 2, 3, 4]]  # 默认预测
                return default_predictions
            return predictions


def create_and_train_enhanced_lstm_crf(csv_file_path, sequence_lstm_model_path, model_save_path=None, lottery_type='plw'):
    """
    创建并训练增强版LSTM-CRF模型
    
    Args:
        csv_file_path: 历史数据文件路径
        sequence_lstm_model_path: 序列LSTM模型文件路径
        model_save_path: 模型保存路径
        lottery_type: 彩票类型 ('plw', 'dlt', 'ssq', 'fc3d')
    """
    print(f"[START] 开始创建和训练{lottery_type}增强版LSTM-CRF模型")
    print("=" * 60)
    
    try:
        # 创建训练器
        trainer = EnhancedLstmCRFTrainer(csv_file_path, model_save_path, lottery_type)
        
        # 1. 加载和准备数据
        X, y = trainer.load_and_prepare_data()
        
        # 2. 训练模式记忆网络
        trainer.train_pattern_memory(y.tolist())
        
        # 3. 加载序列LSTM模型
        trainer.load_sequence_lstm_model(sequence_lstm_model_path)
        
        # 4. 创建增强版LSTM-CRF模型
        trainer.create_enhanced_model()
        
        # 5. 训练模型
        trainer.train_model(X, y, epochs=100)
        
        print(f"[SUCCESS] {lottery_type}增强版LSTM-CRF模型训练完成！")
        return trainer
        
    except Exception as e:
        print(f" 训练过程出错: {e}")
        import traceback
        print(traceback.format_exc())
        return None


# 测试代码
if __name__ == "__main__":
    # 测试增强版LSTM-CRF模型
    print("[TEST] 测试增强版LSTM-CRF模型")
    
    # 创建测试数据
    test_X = np.random.randint(0, 10, (100, 10, 5)).astype(np.float32)
    test_y = np.random.randint(0, 10, (100, 5))
    
    # 创建模型
    model = EnhancedLstmCRFModel(input_dim=5, hidden_dim=64, output_dim=10, output_seq_length=5)
    print("[SUCCESS] 模型创建成功")
    
    # 测试前向传播
    test_input = torch.FloatTensor(test_X[:2])
    with torch.no_grad():
        predictions = model(test_input)
        print(f"[RESULT] 预测结果: {predictions}")